package com.yupi.springbootinit.manager.ai;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import io.reactivex.Scheduler;
import io.reactivex.schedulers.Schedulers;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Component;

import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author admin
 * @date 2025/4/24
 */
@Slf4j
@Component
public class AIManager {

    private static StringBuilder fullContent = new StringBuilder();

    public Flowable<GenerationResult> streamCallWithMessage(Generation gen, List<Message> messageList, ModelEnum model)
            throws NoApiKeyException, ApiException, InputRequiredException {
        GenerationParam param = buildGenerationParam(messageList, model);
        System.out.println("流式输出内容为：");
        Flowable<GenerationResult> result = gen.streamCall(param);
        System.out.println("完整内容为: " + fullContent.toString());
        return result;
    }

    private GenerationParam buildGenerationParam(List<Message> messageList, ModelEnum model) {
        return GenerationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                .model(model.getValue())   // 此处以qwen-plus为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .messages(messageList)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();
    }

    /**
     * 调用接口
     *
     * @param prompt
     * @param model
     * @return
     */
    public Flowable<GenerationResult> callWithMessage(String prompt, ModelEnum model, String systemContent) {
        try {
            Generation gen = new Generation();
            // 用户信息
            Message userMsg = Message.builder()
                    .role(Role.USER.getValue())
                    .content(prompt)
                    .build();
            // 系统信息
            Message systemMsg = Message.builder()
                    .role(Role.SYSTEM.getValue())
                    .content(systemContent)
                    .build();
            List<Message> messageList = Arrays.asList(systemMsg, userMsg);
            return streamCallWithMessage(gen, messageList, model);
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            log.error("An exception occurred: {}", e.getMessage());
        }
        return null;
    }

    /**
     * 默认模型调用 qwen-plus
     */
    public Flowable<GenerationResult> callWithMessage(String prompt, String systemContent) {
        return callWithMessage(prompt, ModelEnum.QWEN_PLUS, systemContent);
    }

    /**
     * 返回String
     */
    public String callWithMessageString(String prompt, String systemContent) {
        return fullContent.toString();
    }


    /**
     * 创建一个线程池，用于执行定时任务
     * @return
     */
    @Bean
    public Scheduler vipScheduler() {
        ThreadFactory threadFactory = new ThreadFactory() {

            final AtomicInteger count = new AtomicInteger(1);

            @Override
            public Thread newThread(@NotNull Runnable r) {
                Thread thread = new Thread(r, "vip-schedule-thread-" + count.getAndIncrement());
                thread.setDaemon(false);
                return thread;
            }
        };

        ExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(10, threadFactory);
        // 返回Scheduler
        return Schedulers.from(scheduledExecutorService);
    }

}
